package com.alibaba.cloud.ai.example.graph.basic_chatbot;

import lombok.NonNull;
import org.bsc.langgraph4j.serializer.plain_text.PlainTextStateSerializer;
import org.bsc.langgraph4j.state.AgentState;
import org.bsc.langgraph4j.state.AgentStateFactory;

import java.io.IOException;
import java.io.ObjectInput;
import java.io.ObjectOutput;
import java.util.HashMap;
import java.util.Map;

public abstract class MemoryStateSerializer<State extends AgentState> extends PlainTextStateSerializer<State> {

    Map<String, Object> data = new HashMap<>();

    protected MemoryStateSerializer(@NonNull AgentStateFactory<State> stateFactory) {
        super(stateFactory);

    }

    @Override
    public String mimeType() {
        return "application/memory";
    }

    @Override
    public void write(State object, ObjectOutput out) throws IOException {
        data.put("state", object);

    }

    @Override
    public State read(ObjectInput in) throws IOException, ClassNotFoundException {
        return (State) data.get("state");
    }
}